package org.cybzacg.encryption.impl;

import org.cybzacg.encryption.core.AbstractEncryptor;
import org.cybzacg.encryption.context.EncryptionContext;
import org.cybzacg.encryption.model.ValidationResult;
import org.cybzacg.encryption.enums.EncryptionAlgorithm;
import org.cybzacg.encryption.exception.EncryptionException;

import javax.crypto.Cipher;
import java.security.KeyPair;
import java.security.KeyPairGenerator;
import java.security.PrivateKey;
import java.security.PublicKey;
import java.security.SecureRandom;
import java.security.spec.PKCS8EncodedKeySpec;
import java.security.spec.X509EncodedKeySpec;
import java.util.Base64;

/**
 * RSA加密器实现
 * 支持RSA-1024、RSA-2048、RSA-4096加密算法
 * 
 * @author encryption team
 * @version 1.0
 */
public class RSAEncryptor extends AbstractEncryptor {

    private static final String ALGORITHM_NAME = "RSA";
    private static final String TRANSFORMATION = "RSA/ECB/PKCS1Padding";

    private volatile KeyPair keyPair;

    /**
     * 构造函数
     */
    public RSAEncryptor() {
        super(EncryptionAlgorithm.RSA);
    }

    /**
     * 构造函数（指定密钥长度）
     * 
     * @param keySize 密钥长度（1024、2048、4096）
     */
    public RSAEncryptor(int keySize) {
        super(EncryptionAlgorithm.RSA);
        metadata.put("keySize", keySize);
    }

    @Override
    protected void doInitialize(EncryptionContext context) {
        try {
            // 获取或生成密钥对
            byte[] publicKeyBytes = context.getPublicKey();
            byte[] privateKeyBytes = context.getPrivateKey();

            if (publicKeyBytes != null && privateKeyBytes != null) {
                // 使用提供的密钥对
                X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(publicKeyBytes);
                PKCS8EncodedKeySpec privKeySpec = new PKCS8EncodedKeySpec(privateKeyBytes);

                java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
                PublicKey publicKey = keyFactory.generatePublic(pubKeySpec);
                PrivateKey privateKey = keyFactory.generatePrivate(privKeySpec);

                keyPair = new KeyPair(publicKey, privateKey);
            } else if (publicKeyBytes != null) {
                // 仅提供公钥（只能加密）
                X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(publicKeyBytes);
                java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
                PublicKey publicKey = keyFactory.generatePublic(pubKeySpec);

                keyPair = new KeyPair(publicKey, null);
            } else {
                // 生成新的密钥对
                KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM_NAME);
                keyGen.initialize(getKeySize(context), new SecureRandom());
                keyPair = keyGen.generateKeyPair();
            }

            // 更新元数据
            metadata.put("transformation", TRANSFORMATION);
            metadata.put("keySize", getKeySize(context));
            metadata.put("hasPublicKey", keyPair.getPublic() != null);
            metadata.put("hasPrivateKey", keyPair.getPrivate() != null);
            metadata.put("initialized", true);

        } catch (Exception e) {
            throw new EncryptionException("Failed to initialize RSA encryptor", e);
        }
    }

    @Override
    protected byte[] doEncrypt(byte[] data, EncryptionContext context) {
        try {
            if (keyPair.getPublic() == null) {
                throw new EncryptionException("No public key available for encryption");
            }

            // 每次加密都创建新的 Cipher 实例（保证线程安全）
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.ENCRYPT_MODE, keyPair.getPublic());
            return cipher.doFinal(data);
        } catch (Exception e) {
            throw new EncryptionException("RSA encryption failed", e);
        }
    }

    @Override
    protected byte[] doDecrypt(byte[] encryptedData, EncryptionContext context) {
        try {
            if (keyPair.getPrivate() == null) {
                throw new EncryptionException("No private key available for decryption");
            }

            // 每次解密都创建新的 Cipher 实例（保证线程安全）
            Cipher cipher = Cipher.getInstance(TRANSFORMATION);
            cipher.init(Cipher.DECRYPT_MODE, keyPair.getPrivate());
            return cipher.doFinal(encryptedData);
        } catch (Exception e) {
            throw new EncryptionException("RSA decryption failed", e);
        }
    }

    @Override
    protected ValidationResult doValidate(EncryptionContext context) {
        try {
            // 验证密钥长度
            Integer keySize = getKeySize(context);
            if (keySize != null && !algorithm.isValidKeySize(keySize)) {
                return ValidationResult.failure(
                        "Invalid RSA key size",
                        "Key size: " + keySize + " bits (supported: 1024, 2048, 4096)",
                        context);
            }

            // 验证公钥
            byte[] publicKeyBytes = context.getPublicKey();
            if (publicKeyBytes != null) {
                try {
                    X509EncodedKeySpec pubKeySpec = new X509EncodedKeySpec(publicKeyBytes);
                    java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
                    keyFactory.generatePublic(pubKeySpec);
                } catch (Exception e) {
                    return ValidationResult.failure("Invalid public key format", e.getMessage(), context);
                }
            }

            // 验证私钥
            byte[] privateKeyBytes = context.getPrivateKey();
            if (privateKeyBytes != null) {
                try {
                    PKCS8EncodedKeySpec privKeySpec = new PKCS8EncodedKeySpec(privateKeyBytes);
                    java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
                    keyFactory.generatePrivate(privKeySpec);
                } catch (Exception e) {
                    return ValidationResult.failure("Invalid private key format", e.getMessage(), context);
                }
            }

            // 检查是否至少有一个密钥
            if (publicKeyBytes == null && privateKeyBytes == null) {
                return ValidationResult.withWarning(
                        "No keys provided, will generate new key pair",
                        "New key pair will be generated",
                        context);
            }

            return ValidationResult.success("RSA validation passed", context);

        } catch (Exception e) {
            return ValidationResult.failure("RSA validation error", e.getMessage(), context);
        }
    }

    @Override
    protected boolean doSupports(EncryptionContext context) {
        if (context == null) {
            return false;
        }

        // 检查算法匹配
        if (context.getAlgorithm() != EncryptionAlgorithm.RSA) {
            return false;
        }

        // 检查加密类型匹配
        if (context.getEncryptionType() != algorithm.getType()) {
            return false;
        }

        // RSA是非对称加密，至少需要一个密钥
        if (context.getPublicKey() == null && context.getPrivateKey() == null) {
            return false;
        }

        return true;
    }

    @Override
    protected void doReset() {
        // Cipher 每次使用时创建新实例，无需重置
        // 只需确保密钥对仍然有效
    }

    @Override
    protected void doClose() {
        keyPair = null;
        metadata.put("initialized", false);
    }

    /**
     * 获取密钥长度
     * 
     * @param context 上下文
     * @return 密钥长度
     */
    private Integer getKeySize(EncryptionContext context) {
        // 优先使用上下文中的密钥长度
        Integer keySize = (Integer) context.getAttribute("keySize");
        if (keySize != null) {
            return keySize;
        }

        // 使用元数据中的密钥长度
        keySize = (Integer) metadata.get("keySize");
        if (keySize != null) {
            return keySize;
        }

        // 使用配置中的密钥长度
        if (context.getConfig() != null) {
            return context.getConfig().getKeySize();
        }

        // 默认使用2048位
        return 2048;
    }

    /**
     * 生成RSA密钥对
     * 
     * @param keySize 密钥长度
     * @return 密钥对
     */
    public static KeyPair generateKeyPair(int keySize) {
        try {
            KeyPairGenerator keyGen = KeyPairGenerator.getInstance(ALGORITHM_NAME);
            keyGen.initialize(keySize, new SecureRandom());
            return keyGen.generateKeyPair();
        } catch (Exception e) {
            throw new EncryptionException("Failed to generate RSA key pair", e);
        }
    }

    /**
     * 获取当前公钥
     * 
     * @return 公钥字节数组，如果未初始化返回null
     */
    public byte[] getCurrentPublicKey() {
        return keyPair != null && keyPair.getPublic() != null ? keyPair.getPublic().getEncoded() : null;
    }

    /**
     * 获取当前私钥
     * 
     * @return 私钥字节数组，如果未初始化返回null
     */
    public byte[] getCurrentPrivateKey() {
        return keyPair != null && keyPair.getPrivate() != null ? keyPair.getPrivate().getEncoded() : null;
    }

    /**
     * 获取公钥的Base64编码字符串
     * 
     * @return Base64编码的公钥
     */
    public String getPublicKeyBase64() {
        byte[] publicKeyBytes = getCurrentPublicKey();
        return publicKeyBytes != null ? Base64.getEncoder().encodeToString(publicKeyBytes) : null;
    }

    /**
     * 获取私钥的Base64编码字符串
     * 
     * @return Base64编码的私钥
     */
    public String getPrivateKeyBase64() {
        byte[] privateKeyBytes = getCurrentPrivateKey();
        return privateKeyBytes != null ? Base64.getEncoder().encodeToString(privateKeyBytes) : null;
    }

    /**
     * 从Base64字符串加载公钥
     * 
     * @param publicKeyBase64 Base64编码的公钥
     * @return 公钥对象
     */
    public static PublicKey loadPublicKeyFromBase64(String publicKeyBase64) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(publicKeyBase64);
            X509EncodedKeySpec spec = new X509EncodedKeySpec(keyBytes);
            java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
            return keyFactory.generatePublic(spec);
        } catch (Exception e) {
            throw new EncryptionException("Failed to load public key from Base64", e);
        }
    }

    /**
     * 从Base64字符串加载私钥
     * 
     * @param privateKeyBase64 Base64编码的私钥
     * @return 私钥对象
     */
    public static PrivateKey loadPrivateKeyFromBase64(String privateKeyBase64) {
        try {
            byte[] keyBytes = Base64.getDecoder().decode(privateKeyBase64);
            PKCS8EncodedKeySpec spec = new PKCS8EncodedKeySpec(keyBytes);
            java.security.KeyFactory keyFactory = java.security.KeyFactory.getInstance(ALGORITHM_NAME);
            return keyFactory.generatePrivate(spec);
        } catch (Exception e) {
            throw new EncryptionException("Failed to load private key from Base64", e);
        }
    }

    /**
     * 检查是否有公钥
     * 
     * @return 是否有公钥
     */
    public boolean hasPublicKey() {
        return keyPair != null && keyPair.getPublic() != null;
    }

    /**
     * 检查是否有私钥
     * 
     * @return 是否有私钥
     */
    public boolean hasPrivateKey() {
        return keyPair != null && keyPair.getPrivate() != null;
    }

    /**
     * 获取最大加密数据长度
     * 
     * @return 最大加密数据长度（字节）
     */
    public int getMaxDataSize() {
        if (keyPair == null || keyPair.getPublic() == null) {
            return 0;
        }

        int keySize = keyPair.getPublic().getEncoded().length * 8;
        // RSA/PKCS1Padding的最大数据长度 = 密钥长度/8 - 11
        return (keySize / 8) - 11;
    }

    @Override
    public String getDescription() {
        return "RSA (Rivest-Shamir-Adleman) asymmetric encryption algorithm";
    }
}
